"""添加MCP配置字段的迁移脚本

这个脚本为MCPTool表添加mcp_config字段，用于存储MCP服务配置
"""

from sqlalchemy import JSON, Column, MetaData, Table, inspect
from sqlalchemy.sql import text

def upgrade(connection):
    """添加mcp_config字段到mcptool表"""
    # 使用无参数初始化MetaData
    meta = MetaData()
    
    # 检查表是否存在
    inspector = inspect(connection)
    if 'mcptool' not in inspector.get_table_names():
        print("mcptool表不存在，无法添加字段")
        return
    
    # 检查列是否已存在
    columns = inspector.get_columns('mcptool')
    column_names = [column['name'] for column in columns]
    
    # 添加列
    if 'mcp_config' not in column_names:
        connection.execute(text('ALTER TABLE mcptool ADD COLUMN mcp_config JSON'))
        print("已添加mcp_config字段到mcptool表")
    else:
        print("mcp_config字段已存在")

def downgrade(connection):
    """移除mcp_config字段"""
    # 使用无参数初始化MetaData
    meta = MetaData()
    
    # 检查表是否存在
    inspector = inspect(connection)
    if 'mcptool' not in inspector.get_table_names():
        print("mcptool表不存在，无需移除字段")
        return
    
    # 检查列是否存在
    columns = inspector.get_columns('mcptool')
    column_names = [column['name'] for column in columns]
    
    # 移除列
    if 'mcp_config' in column_names:
        connection.execute(text('ALTER TABLE mcptool DROP COLUMN mcp_config'))
        print("已从mcptool表移除mcp_config字段")
    else:
        print("mcp_config字段不存在，无需移除") 